﻿#include "CHlsHelper.h"
#include "FileCntl.h"
#include "CHlsToMp4Handler.h"
#include "android_log.h"

#define CHUNK_SIZE 188

static int DoAppend(FileCntl* fpAnterior, FileCntl* fpPosterior, uint64_t fileSize) {
	int ret = 0;
	uint32_t seg = 1024 * 10;
	uint8_t* buffer = new uint8_t[seg];
	uint32_t currentPos = 0;
	do
	{
		uint32_t readSize = (fileSize - currentPos) > seg ? seg : (fileSize - currentPos);
		if (readSize != fpPosterior->readFile(buffer, readSize)) {
			ret = -4;
			break;
		}
		if (readSize != fpAnterior->writeFile(buffer, readSize)) {
			ret = -5;
			break;
		}
		currentPos += readSize;
	} while (currentPos < fileSize);
	delete[] buffer;
	return ret;
}

int CHlsHelper::AppendTs(const char* anteriorSegmentFile, const char* posteriorSegmentFile)
{
	int ret = -1;
	FileCntl* fpAnterior = nullptr;
	FileCntl* fpPosterior = nullptr;
	do 
	{
		ret = -2;
		if (nullptr == anteriorSegmentFile)
		{
			break;
		}
		fpAnterior = new FileCntl(anteriorSegmentFile);
		if (nullptr == fpAnterior) {
			break;
		}
		if (0 != fpAnterior->openFile(O_RDWR | O_CREAT)) {
			break;
		}
		ret = -3;
		if (nullptr == posteriorSegmentFile)
		{
			break;
		}
		fpPosterior = new FileCntl(posteriorSegmentFile);
		if (nullptr == fpPosterior) {
			break;
		}
		if (0 != fpPosterior->openFile(O_RDONLY)) {
			break;
		}
		fpAnterior->seekFile(0, 2);
		uint64_t fileSize = FileCntl::getFileLength(posteriorSegmentFile);
		if (0 == fileSize) {
			break;
		}
		ret = DoAppend(fpAnterior, fpPosterior, fileSize);
	} while (false);
	if (nullptr != fpAnterior)
	{
		fpAnterior->closeFile();
		delete fpAnterior;
		fpAnterior = nullptr;
	}
	if (nullptr != fpPosterior)
	{
		fpPosterior->closeFile();
		delete fpPosterior;
		fpPosterior = nullptr;
	}
	return ret;
}

int CHlsHelper::FileAppend(const char* anteriorFile, const char* posteriorFile) {
	int ret = -1;
	std::fstream* fpAnterior = nullptr;
	std::fstream* fpPosterior = nullptr;
	do
	{
		fpAnterior = new std::fstream(anteriorFile, std::ios::binary | std::ios::in | std::ios::out);
		if (!fpAnterior->is_open()) {
			ret = -2;
			break;
		}
		fpAnterior->seekp(0, std::ios::end);
		fpPosterior = new std::fstream(posteriorFile, std::ios::binary | std::ios::in);
		if (!fpPosterior->is_open()) {
			ret = -3;
			break;
		}
		fpPosterior->seekp(0, std::ios::end);
		unsigned long long fileSize = fpPosterior->tellp();
		fpPosterior->seekp(0, std::ios::beg);
		int tempSize = 1024 * 4;
		char* temp = new char[tempSize];
		bool result = false;
		do
		{
			if (nullptr == temp) {
				ret = -1;
				break;
			}
			int iLength = fileSize > tempSize ? tempSize : fileSize;
			if (!fpPosterior->read(temp, iLength)) {
				ret = -4;
				break;
			}
			if (!fpAnterior->write(temp, iLength)) {
				ret = -5;
				break;
			}
			fileSize -= iLength;
		} while (fileSize > 0);
		if (0 == fileSize) {
			result = true;
			ret = 0;
		}
		if (nullptr != temp) {
			delete[] temp;
			temp = nullptr;
		}
	} while (false);
	if (nullptr != fpAnterior) {
		fpAnterior->close();
		delete fpAnterior;
		fpAnterior = nullptr;
	}
	if (nullptr != fpPosterior) {
		fpPosterior->close();
		delete fpPosterior;
		fpPosterior = nullptr;
	}
	return ret;
}

int CHlsHelper::LocalM3U8ToMp4(const char* m3u8File, const char* outputFile)
{
	int ret = -1;
	do 
	{
		CHlsToMp4Handler handler;
		if (0 != handler.initialize(m3u8File, outputFile)) {
			break;
		}
		handler.run();
		handler.deinitialize();
		ret = 0;
	} while (false);
	return ret;
}

bool CHlsHelper::RenameFile(const char* oldFile, const char* newFile)
{
	int ret = rename(oldFile, newFile);
	return (0 == ret);
}

bool CHlsHelper::FileExist(const char* filePath)
{
	return FileCntl::isFileExist(filePath);
}

bool CHlsHelper::IsTs(const char* filePath)
{
	bool ret = false;
	FileCntl* fp = nullptr;
	do 
	{
		if (nullptr == filePath)
		{
			break;
		}
		fp = new FileCntl(filePath);
		if (nullptr == fp) {
			break;
		}
		if (0 != fp->openFile(O_RDONLY)) {
			break;
		}
		uint8_t flag = 0;
		if (sizeof(flag) != fp->readFile(&flag, sizeof(flag))) {
			break;
		}
		if (flag != 0x47) {
			break;
		}
		if (CHUNK_SIZE != fp->seekFile(CHUNK_SIZE, 0)) {
			break;
		}
		if (sizeof(flag) != fp->readFile(&flag, sizeof(flag))) {
			break;
		}
		if (flag != 0x47) {
			break;
		}
		ret = true;
	} while (false);
	if (fp) {
		fp->closeFile();
		delete fp;
	}
	return ret;
}

bool CHlsHelper::VerifyTsFile(const char* filePath)
{
	bool ret = false;
	do
	{
		if (nullptr == filePath)
		{
			break;
		}
		if (!IsTs(filePath)) {
			break;
		}
		uint64_t length = FileCntl::getFileLength(filePath);
		ret = (length > 0) && (length % CHUNK_SIZE == 0);
	} while (false);
	return ret;
}

static bool IsSameChunk(uint8_t* data1, uint8_t* data2) {
	bool ret = false;
	do 
	{
		int count = 5;
		int segSize = CHUNK_SIZE / count;
		for (int i = 0; i < count; ++i)
		{
			int offset = i * segSize;
			if (offset + sizeof(uint64_t) > CHUNK_SIZE)
			{
				break;
			}
			uint64_t flag1, flag2;
			memcpy(&flag1, data1, sizeof(flag1));
			memcpy(&flag2, data2, sizeof(flag2));
			if (flag1 != flag2) {
				ret = false;
				break;
			}
			else {
				ret = true;
			}
		}
	} while (false);
	return ret;
}

static bool LastChunk(FileCntl* fp, uint8_t* buffer, uint64_t fileSize) {
	bool ret = false;
	do 
	{
		uint64_t seekSize = fileSize - CHUNK_SIZE;
		if (seekSize != fp->seekFile(seekSize, 0)) {
			break;
		}
		
		ret = (fp->readFile(buffer, CHUNK_SIZE) == CHUNK_SIZE);
	} while (false);
	return ret;
}

long long CHlsHelper::FindTsPosition(const char* anteriorFile, const char* posteriorFile)
{
	long long ret = -1;
	FileCntl* fpAnterior = nullptr;
	FileCntl* fpPosterior = nullptr;
	uint8_t* firstAnteriorChunks = nullptr;
	uint8_t anteriorCount = 3;
	uint64_t anteriorChunksLength = CHUNK_SIZE * anteriorCount;
	uint8_t* lastChunk = nullptr;
	uint8_t* tmpChunk = nullptr;
	do 
	{
		if (nullptr == anteriorFile || nullptr == posteriorFile) {
			break;
		}
		if (!IsTs(anteriorFile) || !IsTs(posteriorFile)) {
			break;
		}
		uint64_t anteriorFileLength = FileCntl::getFileLength(anteriorFile);
		uint64_t posteriorFileLength = FileCntl::getFileLength(posteriorFile);
		long long beginPostion = -1;
		uint64_t currentPostion = 0;
		fpAnterior = new FileCntl(anteriorFile);
		if (nullptr == fpAnterior) {
			break;
		}
		fpPosterior = new FileCntl(posteriorFile);
		if (nullptr == fpPosterior) {
			break;
		}
		if (0 != fpAnterior->openFile(O_RDONLY)) {
			break;
		}
		if (0 != fpPosterior->openFile(O_RDONLY)) {
			break;
		}
		firstAnteriorChunks = new uint8_t[anteriorChunksLength];
		if (anteriorChunksLength != fpPosterior->readFile(firstAnteriorChunks, anteriorChunksLength)) {
			break;
		}
		lastChunk = new uint8_t[CHUNK_SIZE];
		if (!LastChunk(fpPosterior, lastChunk, posteriorFileLength)) {
			break;
		}
		tmpChunk = new uint8_t[CHUNK_SIZE];
		bool verifySuccess = false;
		while (currentPostion < anteriorFileLength)
		{
			bool same = false;
			for (int i = 0; i < anteriorCount; ++i)
			{
				same = false;
				if (CHUNK_SIZE != fpAnterior->readFile(tmpChunk, CHUNK_SIZE)) {
					break;
				}
				currentPostion += CHUNK_SIZE;
				if (0 != memcmp(firstAnteriorChunks+i*CHUNK_SIZE, tmpChunk, CHUNK_SIZE))
				{
					break;
				}
				same = true;
			}
			if (!same)
			{
				continue;
			}
			beginPostion = currentPostion - anteriorChunksLength;
			if (anteriorFileLength - currentPostion < posteriorFileLength)
			{
				verifySuccess = true;
				break;
			}
			uint64_t seekSize = beginPostion + (posteriorFileLength - CHUNK_SIZE);
			if (seekSize != fpAnterior->seekFile(seekSize, 0)) {
				break;
			}
			if (CHUNK_SIZE != fpAnterior->readFile(tmpChunk, CHUNK_SIZE)) {
				break;
			}
			if (0 != memcmp(lastChunk, tmpChunk, CHUNK_SIZE))
			{
				if (currentPostion != fpAnterior->seekFile(currentPostion, 0)) {
					break;
				}
				continue;
			}
			currentPostion = seekSize + CHUNK_SIZE;
			verifySuccess = true;
			break;
		}
		if (!verifySuccess)
		{
			break;
		}
		ret = beginPostion;
	} while (false);
	if (fpAnterior) {
		fpAnterior->closeFile();
		delete fpAnterior;
	}
	if (fpPosterior) {
		fpPosterior->closeFile();
		delete fpPosterior;
	}
	if (nullptr != firstAnteriorChunks) {
		delete[] firstAnteriorChunks;
	}
	if (nullptr != lastChunk) {
		delete[] lastChunk;
	}
	if (nullptr != tmpChunk) {
		delete[] tmpChunk;
	}
	return ret;
}


int CHlsHelper::M3U8Decrypt(const char* m3u8File, const char* outputFile)
{
    //LOGI("Input_path=%s, Output_path=%s", in_filename, out_filename);
    AVOutputFormat* ofmt = NULL;
    AVFormatContext* ifmt_ctx = NULL, * ofmt_ctx = NULL;
    AVPacket pkt;
    int ret = -1, i;
    int audioIndex = -1;
    int videoIndex = -1;
    int stream_index = 0;
    int* stream_mapping = NULL;
    int stream_mapping_size;
    int width = 0, height = 0;
    int64_t last_dts = 0;
    int64_t first_pts_dts = 0;
    int first_pkt = 0;
    //int64_t total_packets_count, temp_packets_count = 0;
    //float current_progress, last_progress = -1;

	AVDictionary* opt = nullptr;
	av_dict_set(&opt, "allowed_extensions", "ALL", 0);
	av_dict_set(&opt, "probesize", "2400000", 0);
    if ((ret = avformat_open_input(&ifmt_ctx, m3u8File, 0, &opt)) < 0) {
        avformat_close_input(&ifmt_ctx);
		char strErr[256] = { 0 };
		av_make_error_string(strErr, 256, ret);
		LOGI("m3u8 decrypt open input error:%s",strErr);
        ret = -2;
        return ret;
    }

    if ((ret = avformat_find_stream_info(ifmt_ctx, NULL)) < 0) {
        avformat_close_input(&ifmt_ctx);
        ret = -2;
        return ret;
    }

    av_dump_format(ifmt_ctx, 1, m3u8File, 0);

    //total_packets_count = get_total_packets_count(in_filename);

    ret = avformat_alloc_output_context2(&ofmt_ctx, NULL, "mpegts", outputFile);
    if (!ofmt_ctx) {
        //LOGE("Could not create output context\n");
        char strErr[256] = { 0 };
        av_make_error_string(strErr, 256, ret);
        LOGI("m3u8 decrypt error:%s",strErr);
        avformat_close_input(&ifmt_ctx);
        avformat_free_context(ofmt_ctx);
        ret = -3;
        return ret;
    }
    //LOGI("Output format=%s", ofmt_ctx->oformat->name);

    ofmt_ctx->oformat->flags |= AVFMT_TS_NONSTRICT;
    ofmt_ctx->oformat->flags |= AVFMT_NODIMENSIONS;

    stream_mapping_size = ifmt_ctx->nb_streams;
    stream_mapping = (int*)av_mallocz_array(stream_mapping_size, sizeof(*stream_mapping));
    if (!stream_mapping) {
        //LOGE("Could not alloc stream mapping\n");
        avformat_close_input(&ifmt_ctx);
        avformat_free_context(ofmt_ctx);
        av_freep(&stream_mapping);
		LOGI("m3u8 decrypt error:av_mallocz_array");
        return -1;
    }

    ofmt = ofmt_ctx->oformat;

    for (int i = 0; i < ifmt_ctx->nb_streams; i++) {
        AVStream* out_stream;
        AVStream* in_stream = ifmt_ctx->streams[i];
        AVCodecParameters* in_codecpar = in_stream->codecpar;

        if (in_codecpar->codec_type != AVMEDIA_TYPE_AUDIO &&
            in_codecpar->codec_type != AVMEDIA_TYPE_VIDEO &&
            in_codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE) {
            stream_mapping[i] = -1;
            continue;
        }

        if (in_codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            width = in_codecpar->width;
            height = in_codecpar->height;

            //if (width == 0 && height == 0) {
            //	ret = -1;
            //	return ret;
            //}
            videoIndex = in_stream->index;
        }
        if (in_codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            audioIndex = in_stream->index;
        }

        stream_mapping[i] = stream_index++;

        out_stream = avformat_new_stream(ofmt_ctx, NULL);
        if (!out_stream) {
            ret = -1;
			LOGI("m3u8 decrypt error:avformat_new_stream");
            return ret;
        }

        ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
        if (ret < 0) {
            return ret;
        }
        out_stream->codecpar->codec_tag = 0;
    }
    if (-1 == videoIndex && -1 == audioIndex) {
		LOGI("m3u8 decrypt error:no audio and video");
        ret = -1;
        return ret;
    }

    av_dump_format(ofmt_ctx, 0, outputFile, 1);

    if (!(ofmt->flags & AVFMT_NOFILE)) {
        //LOGI("Open output file");
        ret = avio_open(&ofmt_ctx->pb, outputFile, AVIO_FLAG_WRITE);
        if (ret < 0) {
            /*LOGE("Could not  '%s'", out_filename);
            LOGE("Error occurred: %s\n, width=%d, height=%d", av_err2str(ret), width, height);*/
            avformat_close_input(&ifmt_ctx);
            /* close output */
            if (ofmt_ctx && !(ofmt->flags & AVFMT_NOFILE))
                avio_closep(&ofmt_ctx->pb);
            avformat_free_context(ofmt_ctx);
            av_freep(&stream_mapping);
			LOGI("m3u8 decrypt error:avio_open");
            return ret;
        }
    }

    ret = avformat_write_header(ofmt_ctx, NULL);
    if (ret < 0) {
        //LOGE("Error occurred when opening output file, ret=%d\n", ret);
        //LOGE("Error occurred: %s\n, width=%d, height=%d", av_err2str(ret), width, height);
        avformat_close_input(&ifmt_ctx);
        /* close output */
        if (ofmt_ctx && !(ofmt->flags & AVFMT_NOFILE))
            avio_closep(&ofmt_ctx->pb);
        avformat_free_context(ofmt_ctx);
        av_freep(&stream_mapping);
		LOGI("m3u8 decrypt error:avformat_write_header");
        return ret;
    }

    int64_t video_dts_offset = 0;
    int64_t audio_pts_offset = 0;
    int64_t video_dts_spacing = 0;
    int64_t audio_pts_spacing = 0;
    int64_t last_video_dts = std::numeric_limits<int64_t>::min();
    int64_t last_audio_pts = std::numeric_limits<int64_t>::min();

    while (1) {
        AVStream* in_stream, * out_stream;

        ret = av_read_frame(ifmt_ctx, &pkt);
        if (ret < 0)
            break;

        in_stream = ifmt_ctx->streams[pkt.stream_index];
        if (pkt.stream_index >= stream_mapping_size ||
            stream_mapping[pkt.stream_index] < 0) {
            av_packet_unref(&pkt);
            continue;
        }

        pkt.stream_index = stream_mapping[pkt.stream_index];
        out_stream = ofmt_ctx->streams[pkt.stream_index];

        if (pkt.pts == AV_NOPTS_VALUE) {
            if (pkt.dts != AV_NOPTS_VALUE) {
                pkt.pts = pkt.dts;
                last_dts = pkt.dts;
            }
            else {
                pkt.pts = last_dts + 1;
                pkt.dts = pkt.pts;
                last_dts = pkt.pts;
            }
        }
        else {
            if (pkt.dts != AV_NOPTS_VALUE) {
                last_dts = pkt.dts;
            }
            else {
                pkt.dts = pkt.pts;
                last_dts = pkt.dts;
            }
        }

        if (pkt.pts < pkt.dts) {
            pkt.pts = pkt.dts;
        }

        //起始的时间归0
//        if (!first_pkt && pkt.flags & AV_PKT_FLAG_KEY) {
//            first_pts_dts = pkt.pts;
//            first_pkt++;
//        }

        pkt.pts = pkt.pts - first_pts_dts;
        pkt.dts = pkt.dts - first_pts_dts;
        if (in_stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            if (last_video_dts != std::numeric_limits<int64_t>::min() && 0 == video_dts_spacing) {
                video_dts_spacing = pkt.dts - last_video_dts;
            }
            if (0 != video_dts_spacing) {
                if (pkt.dts < last_video_dts) {
                    int oldDts = pkt.dts;
                    pkt.dts = last_video_dts + video_dts_spacing;
                    pkt.pts = pkt.pts + (pkt.dts - oldDts);
                }
                if (pkt.dts - last_video_dts > 2 * video_dts_spacing) {
                    video_dts_offset += (pkt.dts - last_video_dts - video_dts_spacing);
                }
            }
            last_video_dts = pkt.dts;
            pkt.pts -= video_dts_offset;
            pkt.dts -= video_dts_offset;
        }
        else {
            if (last_audio_pts != std::numeric_limits<int64_t>::min() && 0 == audio_pts_spacing) {
                audio_pts_spacing = pkt.pts - last_audio_pts;
            }
            if (0 != audio_pts_spacing) {
                if (pkt.pts < last_audio_pts) {
                    pkt.pts = last_audio_pts + audio_pts_spacing;
                }
                if (pkt.pts - last_audio_pts > 2 * audio_pts_spacing) {
                    audio_pts_offset += (pkt.pts - last_audio_pts - audio_pts_spacing);
                }
            }
            last_audio_pts = pkt.pts;
            pkt.pts -= audio_pts_offset;
            pkt.dts -= audio_pts_offset;
        }

        /* copy packet */
        pkt.pts = av_rescale_q_rnd(pkt.pts, in_stream->time_base, out_stream->time_base,
                                   static_cast<AVRounding>(AV_ROUND_NEAR_INF |
                                                           AV_ROUND_PASS_MINMAX));
        pkt.dts = av_rescale_q_rnd(pkt.dts, in_stream->time_base, out_stream->time_base,
                                   static_cast<AVRounding>(AV_ROUND_NEAR_INF |
                                                           AV_ROUND_PASS_MINMAX));
        pkt.duration = av_rescale_q(pkt.duration, in_stream->time_base, out_stream->time_base);
        pkt.pos = -1;

        ret = av_interleaved_write_frame(ofmt_ctx, &pkt);

        if (ret < 0) {
            //LOGE("Error muxing packet\n");
            avformat_close_input(&ifmt_ctx);
            /* close output */
            if (ofmt_ctx && !(ofmt->flags & AVFMT_NOFILE))
                avio_closep(&ofmt_ctx->pb);
            avformat_free_context(ofmt_ctx);
            av_freep(&stream_mapping);
            return ret;
        }
        av_packet_unref(&pkt);
    }

    av_write_trailer(ofmt_ctx);

    avformat_close_input(&ifmt_ctx);
    /* close output */
    if (ofmt_ctx && !(ofmt->flags & AVFMT_NOFILE))
        avio_closep(&ofmt_ctx->pb);
    avformat_free_context(ofmt_ctx);
    av_freep(&stream_mapping);
    return 0;
}